from .base import Space
import numpy as np


class Box(Space):
    """
    A box in R^n.
    I.e., each coordinate is bounded.
    """

    def __init__(self, low, high, shape=None):
        """
        初始化盒子的边界。有两种有效的输入方式：

            Box(-1.0, 1.0, (3,4))：low 和 high 是标量，shape 是提供的。
            Box(np.array([-1.0,-2.0]), np.array([2.0,4.0]))：low 和 high 是具有相同形状的数组。

        Two kinds of valid input:

            Box(-1.0, 1.0, (3,4)) # low and high are scalars, and shape is provided
            Box(np.array([-1.0,-2.0]), np.array([2.0,4.0])) # low and high are arrays of the same shape
        """
        if shape is None:
            assert low.shape == high.shape
            self.low = low
            self.high = high
        else:
            assert np.isscalar(low) and np.isscalar(high)
            self.low = low + np.zeros(shape)
            self.high = high + np.zeros(shape)

    def sample(self):
        """从盒子中均匀随机采样一个元素，并返回采样结果"""
        return np.random.uniform(low=self.low, high=self.high, size=self.low.shape)

    def sample_n(self, n):
        """从盒子中均匀随机采样 n 个元素，并返回采样结果。"""
        return np.random.uniform(low=self.low, high=self.high, size=(n,) + self.low.shape)

    def contains(self, x):
        """判断给定的元素 x 是否属于该盒子。判断条件为元素 x 的形状与盒子的形状相同，并且每个坐标轴的取值都在对应的边界范围内。"""
        return x.shape == self.shape and (x >= self.low).all() and (x <= self.high).all()

    @property
    def shape(self):
        """盒子的形状"""
        return self.low.shape

    @property
    def flat_dim(self):
        """扁平化后的向量维度，即将盒子中的元素转换为一维向量后的长度"""
        return int(np.prod(self.low.shape))

    @property
    def bounds(self):
        """返回盒子的边界，即 low 和 high"""
        return self.low, self.high

    def flatten(self, x):
        """将给定的元素进行扁平化"""
        return np.asarray(x).flatten()

    def unflatten(self, x):
        """将给定的元素进行还原操作"""
        return np.asarray(x).reshape(self.shape)

    def flatten_n(self, xs):
        """将给定的元素进行扁平化"""
        xs = np.asarray(xs)
        return xs.reshape((xs.shape[0], -1))

    def unflatten_n(self, xs):
        """将给定的元素进行还原操作"""
        xs = np.asarray(xs)
        return xs.reshape((xs.shape[0],) + self.shape)

    @property
    def default_value(self):
        """返回盒子的默认值，即边界取值的一半"""
        return 0.5 * (self.low + self.high)

    def __repr__(self):
        """返回一个描述盒子形状的字符串"""
        return "Box" + str(self.shape)

    def __eq__(self, other):
        """判断两个盒子是否相等。相等的条件是另一个对象是 Box 类型，并且 low 和 high 的值非常接近"""
        return isinstance(other, Box) and np.allclose(self.low, other.low) and \
               np.allclose(self.high, other.high)

    def __hash__(self):
        """返回盒子对象的哈希值"""
        return hash((self.low, self.high))

